package com.example.springaiollama.controller;

import com.example.springaiollama.service.AiService;
import com.example.springaiollama.service.FunctionCallingService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;

import java.util.Map;

@RestController
@RequestMapping("/api/ai")
@CrossOrigin(origins = "*")
public class AiController {

    private final AiService aiService;
    private final FunctionCallingService functionCallingService;

    @Autowired
    public AiController(AiService aiService, FunctionCallingService functionCallingService) {
        this.aiService = aiService;
        this.functionCallingService = functionCallingService;
    }

    @PostMapping("/chat")
    public ResponseEntity<Map<String, Object>> chat(@RequestBody Map<String, String> request) {
        String userMessage = request.get("message");
        if (userMessage == null || userMessage.trim().isEmpty()) {
            return ResponseEntity.badRequest()
                    .body(Map.of(
                        "success", false,
                        "error", "Message cannot be empty",
                        "timestamp", System.currentTimeMillis()
                    ));
        }

        try {
            String response = aiService.generateResponse(userMessage);
            return ResponseEntity.ok(Map.of(
                "success", true,
                "response", response,
                "timestamp", System.currentTimeMillis()
            ));
        } catch (Exception e) {
            return ResponseEntity.internalServerError()
                    .body(Map.of(
                        "success", false,
                        "error", "Failed to generate response: " + e.getMessage(),
                        "timestamp", System.currentTimeMillis()
                    ));
        }
    }

    @PostMapping("/chat-with-system")
    public ResponseEntity<Map<String, Object>> chatWithSystem(@RequestBody Map<String, String> request) {
        String userMessage = request.get("message");
        String systemPrompt = request.get("systemPrompt");
        
        if (userMessage == null || userMessage.trim().isEmpty()) {
            return ResponseEntity.badRequest()
                    .body(Map.of(
                        "success", false,
                        "error", "Message cannot be empty",
                        "timestamp", System.currentTimeMillis()
                    ));
        }

        try {
            String response;
            if (systemPrompt != null && !systemPrompt.trim().isEmpty()) {
                response = aiService.generateResponseWithSystemPrompt(userMessage, systemPrompt);
            } else {
                response = aiService.generateResponse(userMessage);
            }
            return ResponseEntity.ok(Map.of(
                "success", true,
                "response", response,
                "timestamp", System.currentTimeMillis()
            ));
        } catch (Exception e) {
            return ResponseEntity.internalServerError()
                    .body(Map.of(
                        "success", false,
                        "error", "Failed to generate response: " + e.getMessage(),
                        "timestamp", System.currentTimeMillis()
                    ));
        }
    }

    @PostMapping("/chat-with-function-calling")
    public ResponseEntity<Map<String, Object>> chatWithFunctionCalling(@RequestBody Map<String, String> request) {
        String userMessage = request.get("message");
        if (userMessage == null || userMessage.trim().isEmpty()) {
            return ResponseEntity.badRequest()
                    .body(Map.of(
                        "success", false,
                        "error", "Message cannot be empty",
                        "timestamp", System.currentTimeMillis()
                    ));
        }

        try {
            String response = aiService.generateResponseWithFunctionCalling(userMessage);
            return ResponseEntity.ok(Map.of(
                "success", true,
                "response", response,
                "timestamp", System.currentTimeMillis()
            ));
        } catch (Exception e) {
            return ResponseEntity.internalServerError()
                    .body(Map.of(
                        "success", false,
                        "error", "Failed to generate response: " + e.getMessage(),
                        "timestamp", System.currentTimeMillis()
                    ));
        }
    }

    @PostMapping("/chat-with-function-calling-and-system")
    public ResponseEntity<Map<String, Object>> chatWithFunctionCallingAndSystem(@RequestBody Map<String, String> request) {
        String userMessage = request.get("message");
        String systemPrompt = request.get("systemPrompt");
        
        if (userMessage == null || userMessage.trim().isEmpty()) {
            return ResponseEntity.badRequest()
                    .body(Map.of(
                        "success", false,
                        "error", "Message cannot be empty",
                        "timestamp", System.currentTimeMillis()
                    ));
        }

        try {
            String response;
            if (systemPrompt != null && !systemPrompt.trim().isEmpty()) {
                response = aiService.generateResponseWithFunctionCallingAndSystemPrompt(userMessage, systemPrompt);
            } else {
                response = aiService.generateResponseWithFunctionCalling(userMessage);
            }
            return ResponseEntity.ok(Map.of(
                "success", true,
                "response", response,
                "timestamp", System.currentTimeMillis()
            ));
        } catch (Exception e) {
            return ResponseEntity.internalServerError()
                    .body(Map.of(
                        "success", false,
                        "error", "Failed to generate response: " + e.getMessage(),
                        "timestamp", System.currentTimeMillis()
                    ));
        }
    }

    @GetMapping("/health")
    public ResponseEntity<Map<String, Object>> health() {
        return ResponseEntity.ok(Map.of(
            "success", true,
            "status", "OK",
            "message", "AI Service is running",
            "timestamp", System.currentTimeMillis()
        ));
    }

    @GetMapping("/models")
    public ResponseEntity<Map<String, Object>> getModels() {
        try {
            // 这里可以添加获取可用模型的逻辑
            return ResponseEntity.ok(Map.of(
                "success", true,
                "models", Map.of(
                    "current", "llama2",
                    "available", new String[]{"llama2", "codellama", "mistral", "qwen"}
                ),
                "timestamp", System.currentTimeMillis()
            ));
        } catch (Exception e) {
            return ResponseEntity.internalServerError()
                    .body(Map.of(
                        "success", false,
                        "error", "Failed to get models: " + e.getMessage(),
                        "timestamp", System.currentTimeMillis()
                    ));
        }
    }

    @GetMapping("/tools")
    public ResponseEntity<Map<String, Object>> getAvailableTools() {
        try {
            return ResponseEntity.ok(Map.of(
                "success", true,
                "tools", functionCallingService.getAvailableTools(),
                "timestamp", System.currentTimeMillis()
            ));
        } catch (Exception e) {
            return ResponseEntity.internalServerError()
                    .body(Map.of(
                        "success", false,
                        "error", "Failed to get tools: " + e.getMessage(),
                        "timestamp", System.currentTimeMillis()
                    ));
        }
    }
} 